{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# iflearner client Configure and start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 虚拟环境\n",
    "可以通过conda、pyenv等虚拟工具创建iflearner的虚拟运行环境，然后激活。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1. Conda\n",
    "1.需要安装conda环境。 [安装](https://docs.conda.io/projects/conda/en/stable/user-guide/install/index.html)\n",
    "\n",
    "2.我们将通过conda创建一个虚拟环境并激活它\n",
    "```shell\n",
    "$ conda create -n iflearner python==3.9 ipykernel\n",
    "$ conda activate iflearner\n",
    "```\n",
    "\n",
    "3.将虚拟环境写入jupyter notebook内核\n",
    "命令：python -m ipykernel install --user --name 虚拟环境名称 --display-name 虚拟环境名称\n",
    "- 第一个虚拟环境名称表示创建的虚拟环境名称\n",
    "- 第二个虚拟环境名称表示您希望它出现在 jupyter notebook 的内核选项中\n",
    "\n",
    "   例如： $ python -m ipykernel install --user --name iflearner --display-name \"iflearner\"\n",
    "\n",
    "4.切换内核\n",
    "可以在jupyterlab的内核菜单栏中选择更改内核按钮，选择我们创建的虚拟环境"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.安装iflearner库及相关依赖"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CommandNotFoundError: Your shell has not been properly configured to use 'conda activate'.\n",
      "To initialize your shell, run\n",
      "\n",
      "    $ conda init <SHELL_NAME>\n",
      "\n",
      "Currently supported shells are:\n",
      "  - bash\n",
      "  - fish\n",
      "  - tcsh\n",
      "  - xonsh\n",
      "  - zsh\n",
      "  - powershell\n",
      "\n",
      "See 'conda init --help' for more information and options.\n",
      "\n",
      "IMPORTANT: You may need to close and restart your shell after running 'conda init'.\n",
      "\n",
      "\n",
      "Looking in indexes: http://pypi.douban.com/simple/\n",
      "\u001b[31mERROR: Could not find a version that satisfies the requirement iflearner\u001b[0m\n",
      "\u001b[31mERROR: No matching distribution found for iflearner\u001b[0m\n",
      "Looking in indexes: http://pypi.douban.com/simple/\n",
      "Requirement already satisfied: torch==1.7.1 in /Users/lucky/opt/anaconda3/lib/python3.8/site-packages (1.7.1)\n",
      "Requirement already satisfied: torchvision==0.8.2 in /Users/lucky/opt/anaconda3/lib/python3.8/site-packages (0.8.2)\n",
      "Requirement already satisfied: typing-extensions in /Users/lucky/opt/anaconda3/lib/python3.8/site-packages (from torch==1.7.1) (4.2.0)\n",
      "Requirement already satisfied: numpy in /Users/lucky/opt/anaconda3/lib/python3.8/site-packages (from torch==1.7.1) (1.20.1)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /Users/lucky/opt/anaconda3/lib/python3.8/site-packages (from torchvision==0.8.2) (8.2.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install iflearner torch==1.7.1 torchvision==0.8.2  --index-url http://pypi.douban.com/simple --trusted-host pypi.douban.com "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 定义客户端并启动"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1. 定义 Pytorch 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'iflearner.business'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-deace834deca>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorchvision\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mdatasets\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransforms\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0miflearner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbusiness\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhomo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margument\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mparser\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0miflearner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbusiness\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhomo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpytorch_trainer\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mPyTorchTrainer\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0miflearner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbusiness\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhomo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_client\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mController\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'iflearner.business'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch import nn\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from iflearner.business.homo.argument import parser\n",
    "from iflearner.business.homo.pytorch_trainer import PyTorchTrainer\n",
    "from iflearner.business.homo.train_client import Controller\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, num_channels, num_classes):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(num_channels, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. 集成和实现 PytorchTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mnist(PyTorchTrainer):\n",
    "    def __init__(self, lr=0.15, momentum=0.5) -> None:\n",
    "        self._lr = lr\n",
    "        self._device = (\n",
    "            torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "        )\n",
    "        print(f\"device: {self._device}\")\n",
    "        self._model = Model(num_channels=1, num_classes=10).to(self._device)\n",
    "\n",
    "        super().__init__(self._model)\n",
    "\n",
    "        self._optimizer = optim.SGD(self._model.parameters(), lr=lr, momentum=momentum)\n",
    "        self._loss = F.nll_loss\n",
    "\n",
    "        apply_transform = transforms.Compose(\n",
    "            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
    "        )\n",
    "        train_dataset = datasets.MNIST(\n",
    "            \"./data\", train=True, download=True, transform=apply_transform\n",
    "        )\n",
    "        test_dataset = datasets.MNIST(\n",
    "            \"./data\", train=False, download=True, transform=apply_transform\n",
    "        )\n",
    "        self._train_data = torch.utils.data.DataLoader(\n",
    "            train_dataset, batch_size=64, shuffle=True\n",
    "        )\n",
    "        self._test_data = torch.utils.data.DataLoader(\n",
    "            test_dataset, batch_size=64, shuffle=False\n",
    "        )\n",
    "\n",
    "    def fit(self, epoch):\n",
    "        self._model.to(self._device)\n",
    "        self._model.train()\n",
    "        print(\n",
    "            f\"Epoch: {epoch}, the size of training dataset: {len(self._train_data.dataset)}, batch size: {len(self._train_data)}\"\n",
    "        )\n",
    "        for batch_idx, (data, target) in enumerate(self._train_data):\n",
    "            data, target = data.to(self._device), target.to(self._device)\n",
    "            self._optimizer.zero_grad()\n",
    "            output = self._model(data)\n",
    "            loss = self._loss(output, target)\n",
    "            loss.backward()\n",
    "            self._optimizer.step()\n",
    "\n",
    "    def evaluate(self, epoch):\n",
    "        self._model.to(self._device)\n",
    "        self._model.eval()\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        print(f\"The size of testing dataset: {len(self._test_data.dataset)}\")\n",
    "        with torch.no_grad():\n",
    "            for data, target in self._test_data:\n",
    "                data, target = data.to(self._device), target.to(self._device)\n",
    "                output = self._model(data)\n",
    "                test_loss += self._loss(\n",
    "                    output, target, reduction=\"sum\"\n",
    "                ).item()  # sum up batch loss\n",
    "                pred = output.argmax(\n",
    "                    dim=1, keepdim=True\n",
    "                )  # get the index of the max log-probability\n",
    "                correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "        test_loss /= len(self._test_data.dataset)\n",
    "\n",
    "        print(\n",
    "            \"Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\".format(\n",
    "                test_loss,\n",
    "                correct,\n",
    "                len(self._test_data.dataset),\n",
    "                100.0 * correct / len(self._test_data.dataset),\n",
    "            )\n",
    "        )\n",
    "\n",
    "        return {\"loss\": test_loss, \"acc\": correct / len(self._test_data.dataset)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3. 启动client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    args = parser.parse_args(args=[])\n",
    "    print(args)\n",
    "    args.name = \"client02\"\n",
    "    args.epochs = 2\n",
    "    mnist = Mnist()\n",
    "    controller = Controller(args, mnist)\n",
    "    controller.run()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
